"""
优化器访问者

实现AST的优化功能，用于优化工作流的性能和结构。
"""

from typing import Any, Dict, List, Optional, Set
from .base import BaseAstVisitor

# 前向声明，避免循环导入
from typing import TYPE_CHECKING
if TYPE_CHECKING:
    from ..ast.base import Ast, AstContext
    from ..ast.nodes import *


class OptimizationResult:
    """优化结果"""
    
    def __init__(self):
        self.optimized_ast: Optional['Ast'] = None
        self.optimizations_applied: List[str] = []
        self.performance_improvement: float = 0.0
        self.nodes_removed: int = 0
        self.nodes_merged: int = 0
        self.nodes_reordered: int = 0
        self.warnings: List[str] = []
        self.metadata: Dict[str, Any] = {}
    
    def add_optimization(self, description: str) -> None:
        """添加优化记录"""
        self.optimizations_applied.append(description)
    
    def add_warning(self, warning: str) -> None:
        """添加警告"""
        self.warnings.append(warning)


class OptimizerVisitor(BaseAstVisitor):
    """
    优化器访问者
    
    负责优化AST结构和性能，包括：
    - 死代码消除
    - 节点合并
    - 执行顺序优化
    - 并行度优化
    - 资源使用优化
    """
    
    def __init__(self):
        super().__init__()
        self.optimization_result = OptimizationResult()
        self.optimization_level = 1  # 0=无优化, 1=基础优化, 2=激进优化
        self.preserve_semantics = True  # 是否保持语义不变
    
    def optimize(self, ast: 'Ast', context: 'AstContext') -> OptimizationResult:
        """优化AST"""
        # 重置优化结果
        self.optimization_result = OptimizationResult()
        
        # 执行优化
        optimized_ast = self.visit(ast, context)
        
        # 设置优化结果
        self.optimization_result.optimized_ast = optimized_ast
        
        return self.optimization_result
    
    def visit(self, ast: 'Ast', context: 'AstContext') -> 'Ast':
        """重写visit方法，添加优化逻辑"""
        # 应用节点级优化
        optimized_ast = self._optimize_node(ast, context)
        
        # 调用父类方法处理子节点
        result = super().visit(optimized_ast, context)
        
        # 应用后处理优化
        result = self._post_optimize_node(result, context)
        
        return result
    
    def _optimize_node(self, ast: 'Ast', context: 'AstContext') -> 'Ast':
        """优化单个节点"""
        # 克隆节点以避免修改原始AST
        optimized = ast.clone()
        
        # 应用各种优化策略
        optimized = self._eliminate_dead_code(optimized, context)
        optimized = self._optimize_constants(optimized, context)
        optimized = self._optimize_conditions(optimized, context)
        
        return optimized
    
    def _post_optimize_node(self, ast: 'Ast', context: 'AstContext') -> 'Ast':
        """后处理优化"""
        # 应用结构优化
        optimized = self._merge_adjacent_nodes(ast, context)
        optimized = self._reorder_nodes(optimized, context)
        
        return optimized
    
    def _eliminate_dead_code(self, ast: 'Ast', context: 'AstContext') -> 'Ast':
        """消除死代码"""
        # 检查条件节点的死分支
        if ast.node_type.value == 'conditional':
            condition = getattr(ast, 'condition', '')
            
            # 如果条件是常量，可以消除一个分支
            if condition.lower() == 'true':
                if hasattr(ast, 'true_branch') and ast.true_branch:
                    self.optimization_result.add_optimization("消除条件节点的false分支")
                    return ast.true_branch
            elif condition.lower() == 'false':
                if hasattr(ast, 'false_branch') and ast.false_branch:
                    self.optimization_result.add_optimization("消除条件节点的true分支")
                    return ast.false_branch
        
        # 检查空的序列节点
        if ast.node_type.value == 'sequence':
            if hasattr(ast, 'children') and not ast.children:
                self.optimization_result.add_optimization("消除空的序列节点")
                self.optimization_result.nodes_removed += 1
                return None
        
        return ast
    
    def _optimize_constants(self, ast: 'Ast', context: 'AstContext') -> 'Ast':
        """优化常量"""
        # 解析常量表达式
        if hasattr(ast, 'max_duration') and isinstance(ast.max_duration, str):
            try:
                # 尝试将字符串转换为数字
                ast.max_duration = float(ast.max_duration)
                self.optimization_result.add_optimization("优化常量表达式")
            except ValueError:
                pass
        
        return ast
    
    def _optimize_conditions(self, ast: 'Ast', context: 'AstContext') -> 'Ast':
        """优化条件表达式"""
        if ast.node_type.value == 'conditional':
            condition = getattr(ast, 'condition', '')
            
            # 简化条件表达式
            if condition == '1' or condition.lower() == 'yes':
                ast.condition = 'true'
                self.optimization_result.add_optimization("简化条件表达式")
            elif condition == '0' or condition.lower() == 'no':
                ast.condition = 'false'
                self.optimization_result.add_optimization("简化条件表达式")
        
        return ast
    
    def _merge_adjacent_nodes(self, ast: 'Ast', context: 'AstContext') -> 'Ast':
        """合并相邻的相同类型节点"""
        if not hasattr(ast, 'children') or not ast.children:
            return ast
        
        # 合并相邻的序列节点
        if ast.node_type.value == 'sequence':
            merged_children = []
            i = 0
            while i < len(ast.children):
                child = ast.children[i]
                
                # 如果子节点也是序列节点，展开它
                if child.node_type.value == 'sequence' and hasattr(child, 'children'):
                    merged_children.extend(child.children)
                    self.optimization_result.add_optimization("展开嵌套的序列节点")
                    self.optimization_result.nodes_merged += 1
                else:
                    merged_children.append(child)
                
                i += 1
            
            ast.children = merged_children
        
        # 合并相邻的文件操作节点
        elif ast.node_type.value == 'sequence':
            merged_children = []
            i = 0
            while i < len(ast.children):
                child = ast.children[i]
                
                # 检查是否可以与下一个节点合并
                if (i + 1 < len(ast.children) and 
                    child.node_type.value == 'file_operation' and
                    ast.children[i + 1].node_type.value == 'file_operation'):
                    
                    # 简单的合并逻辑（实际实现会更复杂）
                    if (hasattr(child, 'operation') and child.operation == 'copy' and
                        hasattr(ast.children[i + 1], 'operation') and ast.children[i + 1].operation == 'delete'):
                        
                        # 将copy+delete合并为move
                        merged_node = child.clone()
                        merged_node.operation = 'move'
                        merged_children.append(merged_node)
                        
                        self.optimization_result.add_optimization("合并文件操作: copy+delete -> move")
                        self.optimization_result.nodes_merged += 1
                        i += 2  # 跳过下一个节点
                        continue
                
                merged_children.append(child)
                i += 1
            
            ast.children = merged_children
        
        return ast
    
    def _reorder_nodes(self, ast: 'Ast', context: 'AstContext') -> 'Ast':
        """重新排序节点以优化性能"""
        if not hasattr(ast, 'children') or not ast.children:
            return ast
        
        # 对于序列节点，将快速节点放在前面
        if ast.node_type.value == 'sequence':
            # 定义节点执行时间权重
            time_weights = {
                'scan_directory': 1,
                'extract_metadata': 2,
                'data_transform': 1,
                'file_operation': 2,
                'organize_files': 3,
                'ffmpeg_slice': 10,
                'gemini_classify': 8,
                'output_draft': 5,
                'generate_report': 3,
            }
            
            # 按执行时间排序（快的在前）
            original_order = list(range(len(ast.children)))
            ast.children.sort(key=lambda node: time_weights.get(node.node_type.value, 5))
            new_order = [ast.children.index(child) for child in ast.children]
            
            if original_order != new_order:
                self.optimization_result.add_optimization("重新排序节点以优化执行顺序")
                self.optimization_result.nodes_reordered += len(ast.children)
        
        return ast
    
    # ========================================================================
    # 具体节点的优化实现
    # ========================================================================
    
    def visit_sequence(self, ast: 'Sequence', context: 'AstContext') -> 'Sequence':
        """优化顺序执行节点"""
        # 优化子节点
        optimized_children = []
        for child in ast.children:
            optimized_child = self.visit(child, context)
            if optimized_child is not None:  # 过滤掉被消除的节点
                optimized_children.append(optimized_child)
        
        # 如果只有一个子节点，可以消除序列节点
        if len(optimized_children) == 1:
            self.optimization_result.add_optimization("消除单子节点的序列节点")
            self.optimization_result.nodes_removed += 1
            return optimized_children[0]
        
        # 如果没有子节点，返回None（将被消除）
        if not optimized_children:
            self.optimization_result.add_optimization("消除空的序列节点")
            self.optimization_result.nodes_removed += 1
            return None
        
        optimized_ast = ast.clone()
        optimized_ast.children = optimized_children
        return optimized_ast
    
    def visit_parallel(self, ast: 'Parallel', context: 'AstContext') -> 'Parallel':
        """优化并行执行节点"""
        # 优化子节点
        optimized_children = []
        for child in ast.children:
            optimized_child = self.visit(child, context)
            if optimized_child is not None:
                optimized_children.append(optimized_child)
        
        # 优化并行度
        optimized_ast = ast.clone()
        optimized_ast.children = optimized_children
        
        # 根据子节点数量调整并行度
        if len(optimized_children) < ast.max_workers:
            optimized_ast.max_workers = len(optimized_children)
            self.optimization_result.add_optimization(
                f"调整并行度从 {ast.max_workers} 到 {optimized_ast.max_workers}"
            )
        
        # 如果只有一个子节点，转换为序列节点
        if len(optimized_children) == 1:
            from ..ast.nodes import Sequence
            sequence_node = Sequence(children=optimized_children)
            sequence_node.node_id = ast.node_id
            sequence_node.name = ast.name
            sequence_node.description = ast.description
            
            self.optimization_result.add_optimization("将单子节点的并行节点转换为序列节点")
            return sequence_node
        
        return optimized_ast
    
    def visit_conditional(self, ast: 'Conditional', context: 'AstContext') -> 'Conditional':
        """优化条件执行节点"""
        optimized_ast = ast.clone()
        
        # 优化分支
        if ast.true_branch:
            optimized_true = self.visit(ast.true_branch, context)
            optimized_ast.true_branch = optimized_true
        
        if ast.false_branch:
            optimized_false = self.visit(ast.false_branch, context)
            optimized_ast.false_branch = optimized_false
        
        # 如果两个分支都为空，消除条件节点
        if not optimized_ast.true_branch and not optimized_ast.false_branch:
            self.optimization_result.add_optimization("消除空分支的条件节点")
            self.optimization_result.nodes_removed += 1
            return None
        
        return optimized_ast
    
    def visit_loop(self, ast: 'Loop', context: 'AstContext') -> 'Loop':
        """优化循环执行节点"""
        # 检查循环展开的可能性
        if (hasattr(ast, 'max_iterations') and 
            ast.max_iterations and 
            ast.max_iterations <= 3 and
            self.optimization_level >= 2):
            
            # 小循环可以展开
            from ..ast.nodes import Sequence
            
            expanded_children = []
            for i in range(ast.max_iterations):
                for child in ast.children:
                    expanded_child = child.clone()
                    expanded_children.append(expanded_child)
            
            sequence_node = Sequence(children=expanded_children)
            sequence_node.node_id = ast.node_id
            sequence_node.name = f"{ast.name}_expanded"
            sequence_node.description = f"展开的循环: {ast.description}"
            
            self.optimization_result.add_optimization(
                f"展开小循环 (迭代次数: {ast.max_iterations})"
            )
            
            return self.visit(sequence_node, context)
        
        # 优化子节点
        optimized_children = []
        for child in ast.children:
            optimized_child = self.visit(child, context)
            if optimized_child is not None:
                optimized_children.append(optimized_child)
        
        optimized_ast = ast.clone()
        optimized_ast.children = optimized_children
        return optimized_ast
    
    def visit_ffmpeg_slice(self, ast: 'FfmpegSlice', context: 'AstContext') -> 'FfmpegSlice':
        """优化FFmpeg切片节点"""
        optimized_ast = ast.clone()
        
        # 优化质量设置
        if ast.quality == 'high' and self.optimization_level == 0:
            optimized_ast.quality = 'medium'
            self.optimization_result.add_optimization("降低视频质量以提高处理速度")
        
        # 优化滤镜
        if hasattr(ast, 'filters') and ast.filters:
            # 移除重复的滤镜
            unique_filters = list(dict.fromkeys(ast.filters))
            if len(unique_filters) < len(ast.filters):
                optimized_ast.filters = unique_filters
                self.optimization_result.add_optimization("移除重复的视频滤镜")
        
        return optimized_ast
    
    def visit_gemini_classify(self, ast: 'GeminiClassify', context: 'AstContext') -> 'GeminiClassify':
        """优化Gemini分类节点"""
        optimized_ast = ast.clone()
        
        # 优化温度参数
        if ast.temperature > 0.5:
            optimized_ast.temperature = 0.1
            self.optimization_result.add_optimization("降低AI模型温度以提高一致性")
        
        # 优化token数量
        if ast.max_tokens > 500:
            optimized_ast.max_tokens = 500
            self.optimization_result.add_optimization("减少最大token数量以提高响应速度")
        
        return optimized_ast
    
    # 其他节点使用默认实现
    def visit_default(self, ast: 'Ast', context: 'AstContext') -> 'Ast':
        """默认优化方法"""
        return ast.clone()
